"""
FEniCS tutorial demo program: Poisson equation with Dirichlet conditions.
Test problem is chosen to give an exact solution at all nodes of the mesh.

  -Laplace(u) = f    in the unit square
            u = u_D  on the boundary

  u_D = 1 + x^2 + 2y^2
    f = -6


ref: https://www.iesensor.com/blog/2017/05/24/gmsh_fenics_meshing/
     https://comphysblog.wordpress.com/2018/08/15/fenics-2d-electrostatics-with-imported-mesh-and-boundaries/
"""

from __future__ import print_function
from fenics import *
from dolfin import *
import matplotlib.pyplot as plt
import os
#from vtkplotter.dolfin import plot
import numpy as np
import math

class PlasmaHModel:
    def __init__(self):
        pass

    def run(self):
        #fname = "t1"
        fname = "test_poisson_hl2a"
        mesh = Mesh(fname+".xml")
        if os.path.exists( fname+"_physical_region.xml"):
            subdomains = MeshFunction("size_t", mesh, fname+"_physical_region.xml")
            #plot(subdomains)
        if os.path.exists( fname+"_facet_region.xml"):
            boundaries = MeshFunction("size_t", mesh, fname+"_facet_region.xml")
            #plot(boundaries)

        #plot(mesh)
        #plt.show()


        print("xml mesh reading done")


        V = FunctionSpace(mesh, 'P', 1)



        # Define boundary condition
        u_D = Expression('1 + x[0]*x[0] + 2*x[1]*x[1]', degree=2)

        def boundary(x, on_boundary):
            return on_boundary

        #bc = DirichletBC(V, u_D, boundary)

        #the parameter after boundaries, see hl2a_facet_region.xml file
        inner_edge_boundary = DirichletBC(V, Constant(100.0), boundaries, 384)
        #first_wall_boundary = DirichletBC(V, Constant(50.0), boundaries, 385)
        divertor_ll_boundary = DirichletBC(V, Constant(10.0), boundaries, 386)
        divertor_lr_boundary = DirichletBC(V, Constant(10.0), boundaries,387)
        #dome_l_boundary = DirichletBC(V, Constant(10.0), boundaries, 388)

        bcs =[inner_edge_boundary, divertor_ll_boundary, divertor_lr_boundary]
        #bcs =[inner_edge_boundary, first_wall_boundary, divertor_ll_boundary, divertor_lr_boundary, dome_l_boundary]





        # Define variational problem
        u = TrialFunction(V)
        v = TestFunction(V)
        f = Constant(10.0)
        a = dot(grad(u), grad(v))*dx
        L = f*v*dx

        # Compute solution
        u = Function(V)
        solve(a == L, u, bcs)

        # Plot solution and mesh
        c = plot(u, mode='color')
        plt.colorbar(c)

        #plot(mesh)

        #plot using vtk
        #plot(u, mode='color', style=1)


        # Save solution to file in VTK format
        #vtkfile = File('poisson/solution.pvd')
        #vtkfile << u

        # Compute error in L2 norm
        error_L2 = errornorm(u_D, u, 'L2')

        # Compute maximum error at vertices
        vertex_values_u_D = u_D.compute_vertex_values(mesh)
        vertex_values_u = u.compute_vertex_values(mesh)
        import numpy as np
        error_max = np.max(np.abs(vertex_values_u_D - vertex_values_u))

        # Print errors
        print('error_L2  =', error_L2)
        print('error_max =', error_max)

        # Hold plot
        plt.show()
